from scipy import stats
import numpy as np


def extract_features(signal, fs=1000):
    """
    参数说明：
    signal : 输入信号（时域波形，推荐长度≥1000）
    fs     : 采样率
    
    返回：特征字典（包含峰度、RMS、峰值频率）
    """
    # features = {}
    #
    # # 1. 时域特征
    # features['rms'] = np.sqrt(np.mean(signal ** 2))  # 有效值
    #
    # # 2. 峰度（原始定义，高斯分布为3.0）
    # features['kurtosis'] = stats.kurtosis(signal, fisher=False)
    #
    # # 3. 频域特征
    # fft = np.abs(np.fft.fft(signal))
    # freqs = np.fft.fftfreq(len(signal), 1 / fs)
    # main_freq = freqs[np.argmax(fft)]
    # features['peak_freq'] = np.abs(main_freq)  # 取主频绝对值
    #
    # return features
    # features = {}
    # # 原始特征提取
    # features['rms'] = np.sqrt(np.mean(signal ** 2))
    # features['kurtosis'] = stats.kurtosis(signal, fisher=False)
    # fft = np.abs(np.fft.fft(signal))
    # freqs = np.fft.fftfreq(len(signal), 1 / fs)
    # features['peak_freq'] = np.abs(freqs[np.argmax(fft)])
    #
    # # 计算综合评分（示例：加权平均）
    # composite_score = 0.5 * features['rms'] + 0.3 * features['kurtosis'] + 0.2 * features['peak_freq']
    # return {'composite': composite_score}  # 返回单一特征
    """返回包含所有特征的字典"""
    features = {}

    # 1. 时域特征
    features['rms'] = np.sqrt(np.mean(signal ** 2))  # 有效值

    # 2. 峰度（原始定义，高斯分布为3.0）
    features['kurtosis'] = stats.kurtosis(signal, fisher=False)

    # 3. 频域特征
    fft = np.abs(np.fft.fft(signal))
    freqs = np.fft.fftfreq(len(signal), 1 / fs)
    features['peak_freq'] = np.abs(freqs[np.argmax(fft)])

    # 4. 综合评分（示例：加权平均）
    features['composite'] = 0.5 * features['rms'] + 0.3 * features['kurtosis'] + 0.2 * features['peak_freq']

    return features  # 返回所有特征
